import numpy as np
import pygame


class SHAPE(object):

    def __init__(self, grid_num):
        self.npl = grid_num[0]  # 列数
        self.nph = grid_num[1]  # 行数
        self.grid_np = np.zeros((self.nph, self.npl), dtype='int8')

        self.shape_np = np.array([[[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 1],  #oooo
                                  [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 4],
                                  [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 1],
                                  [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 4]],

                                 [[1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2],  #oo
                                  [1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2],  #oo
                                  [1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2],
                                  [1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2]],

                                  [[0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],  # o
                                   [1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 3],  #ooo
                                   [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],
                                   [0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 3]],

                                  [[1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],  # oo
                                   [0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 3],  #  oo
                                   [1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],
                                   [0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 3]],

                                  [[0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],  #  oo
                                   [1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 3],  # oo
                                   [0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],
                                   [1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 3]],

                                  [[1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],  # ooo
                                   [1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 2, 3],  # o
                                   [0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],
                                   [1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 2, 3]],

                                  [[1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],  # ooo
                                   [0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 2, 3],  #   o
                                   [1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 2],
                                   [1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 3]]
                                  ], dtype='int8')
        self.shape_vertex = [int(self.npl/2), 0]          # 列、行
        self.shape_type_num = self.shape_np.shape[0]
        self.shape_rota_num = self.shape_np.shape[1]
        self.shape_type_index = np.random.randint(0, self.shape_type_num)
        self.shape_rota_index = np.random.randint(0, self.shape_rota_num)

        self.shape_color = pygame.Color(145, 145, 145)
        self.grid_color = pygame.Color(255, 255, 255)

        self.whether_cut = 0
        self.whether_restart = 0
        self.whether_merge = 0
        self.cavity_sum = 0

    def reset_shape(self):
        """重置形状"""
        self.grid_np = np.zeros((self.nph, self.npl))
        self.shape_vertex = [int(self.npl / 2), 0]
        self.shape_type_index = np.random.randint(0, self.shape_type_num)
        self.shape_rota_index = np.random.randint(0, self.shape_rota_num)

    def chiose_shape(self):
        """选择形状,如果一出来就撞到其它，说明游戏结束"""
        self.shape_vertex = [int(self.npl / 2), 0]
        self.shape_type_index = np.random.randint(0, self.shape_type_num)
        self.shape_rota_index = np.random.randint(0, self.shape_rota_num)
        if self.check_hill_other(self.shape_vertex[0], self.shape_vertex[1]) == 1:
            self.reset_shape()
            self.whether_restart = 1

    def rota_shape(self):
        """转动形状"""
        # 转动序号加一.不能大于等于4，需要回到0
        self.shape_rota_index += 1
        if self.shape_rota_index == 4:
            self.shape_rota_index = 0
        # 如果旋转会撞到墙或者撞到其它方块，则不旋转
        if self.check_hill_wo(self.shape_vertex[0]) == 1:
            self.shape_rota_index -= 1

    def move_step(self, action):
        """往下掉、左边、右边一步"""
        l, h = self.shape_vertex
        if action == 'down':
            if self.check_hill_bo(h+1) == 0:
                h += 1
            else:
                self.merge_shape_grid()
                self.cut_h()
                self.chiose_shape()
                return
        elif action == 'right':
            if self.check_hill_wo(l+1) == 0:
                l += 1
        elif action == 'left':
            if self.check_hill_wo(l-1) == 0:
                l -= 1
        self.shape_vertex = [l, h]

    def check_hill_other(self, l, h):
        #  hill the other
        shape_data = self.shape_np[self.shape_type_index][self.shape_rota_index]
        shape_data = shape_data[0:-2].reshape((4, 4))
        x = np.where(shape_data == 1)[0] + h
        y = np.where(shape_data == 1)[1] + l
        if np.max(x) >= self.nph:
            return 1
        else:
            if np.sum(self.grid_np[x, y]) >= 1:
                return 1
            else:
                return 0

    def check_hill_wo(self, l):
        """检查是否左右撞墙"""
        # hill the wall
        sl = self.shape_np[self.shape_type_index][self.shape_rota_index][-2]
        if l < 0 or (l + sl) > self.npl:
            return 1
        # hill other
        else:
            return self.check_hill_other(l, self.shape_vertex[1])

    def check_hill_bo(self, h):
        """检查是否撞到底部的东西，需要合并"""
        # hill the bottom
        sh = self.shape_np[self.shape_type_index][self.shape_rota_index][-1]
        if (h + sh) > self.nph:
            return 1
        # hill other
        else:
            return self.check_hill_other(self.shape_vertex[0], h)

    def merge_shape_grid(self):
        """合并形状和棋盘"""
        shape_data = self.shape_np[self.shape_type_index][self.shape_rota_index]
        shape_data = shape_data[0:-2].reshape((4, 4))
        # 计算将会产生多少空洞
        self.whether_merge = 1
        x = np.where(shape_data == 1)[0] + self.shape_vertex[1]
        y = np.where(shape_data == 1)[1] + self.shape_vertex[0]
        self.cavity_sum = 0
        for index in set(y):
            if len(np.where(self.grid_np[:, index] == 1)[0]) != 0:
                old_height = np.where(self.grid_np[:, index] == 1)[0].min()
            else:
                old_height = self.nph
            new_height = x[np.where(y == index)[0]].max()
            self.cavity_sum += (old_height - new_height - 1)
        # 合并数据
        self.grid_np[x, y] = 1


    def cut_h(self):
        """一行满，则该行消失"""
        lm_sum = np.sum(self.grid_np, axis=1)
        cut_position = np.where(lm_sum == self.npl)
        if len(cut_position[0]) > 0:
            self.whether_cut = 1
            for index in cut_position[0]:
                self.grid_np = np.delete(self.grid_np, index, axis=0)
                self.grid_np = np.insert(self.grid_np, 0, np.zeros(self.npl), axis=0)

    def draw_shape(self, window, grid_size):
        """画形状和棋盘"""
        # 画
        for zl in range(self.npl):
            for zh in range(self.nph):
                if self.grid_np[zh][zl] == 1:
                    x = zl * grid_size[0]
                    y = zh * grid_size[1]
                    rec = (x+1, y+1, grid_size[0]-2, grid_size[1]-2)
                    pygame.draw.rect(window, self.grid_color, rec)
        # 画
        shape_data = self.shape_np[self.shape_type_index][self.shape_rota_index]
        for index, data in enumerate(shape_data):
            if index <= 15 and data == 1:
                x = ((index % 4) + self.shape_vertex[0]) * grid_size[0]  # 列
                y = (int(index / 4) + self.shape_vertex[1]) * grid_size[1]
                rec = (x+1, y+1, grid_size[0]-2, grid_size[1]-2)
                pygame.draw.rect(window, self.shape_color, rec)

    def get_s(self):
        """计算状态"""
        s1 = self.grid_np.reshape(1, self.npl * self.nph)
        s2 = np.array([self.shape_vertex[0], self.shape_vertex[1], self.shape_type_index, self.shape_rota_index])
        return np.hstack((s1[0], s2))

    def get_r(self, r):
        """计算奖励"""
        if self.whether_cut == 1:
            self.whether_cut = 0
            r += 1
        if self.whether_restart == 1:
            self.whether_restart = 0
            return 0
        # 投影下去

        # 导致产生了空洞,整体高度增加
        if self.whether_merge == 1:
            self.whether_merge = 0
            r -= self.cavity_sum

            max_h = np.sum(self.grid_np, axis=0).max()
            r -= max_h
        return r
